/*
 * SqueamishOssifrage.c
 * CSAW 2011 CTF Challenge Kernel Module
 * Jon Oberheide <jon@oberheide.org>
 *
 * Decrypt the flag...and win?
 *
 * Information:
 *
 *   /proc/csaw/key:     get/set symmetric key
 *   /proc/csaw/encrypt: get/set encrypted payload
 *   /proc/csaw/decrypt: get/set decrypted payload
 *
 * Example usage:
 *
 *   # set a key
 *   $ echo "this is a key" > /proc/csaw/key
 *
 *   # encrypt payload with that key
 *   $ echo "hello world" > /proc/csaw/encrypt
 *
 *   # retrieve encrypted payload
 *   $ cat /proc/csaw/encrypt > payload.enc
 *
 *   # decrypt encrypted payload
 *   $ cat payload.enc > /proc/csaw/decrypt
 *
 *   # retrieve decrypted payload
 *   $ cat /proc/csaw/decrypt > payload.dec
 *
 *   # view decrypted payload
 *   $ cat payload.dec
 *   hello world
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/proc_fs.h>
#include <linux/string.h>
#include <asm/uaccess.h>

#define PROC_CSAW    "csaw2011"
#define PROC_KEY     "key"
#define PROC_ENCRYPT "encrypt"
#define PROC_DECRYPT "decrypt"
#define MAX_LENGTH   128
#define MAX_ROUNDS   64
#define DELTA        0x9e3779b9

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Jon Oberheide");
MODULE_DESCRIPTION("CSAW 2011 CTF Challenge Kernel Module");

static struct proc_dir_entry *proc_csaw;
static struct proc_dir_entry *proc_key;
static struct proc_dir_entry *proc_decrypt;
static struct proc_dir_entry *proc_encrypt;

static uint32_t key[4] = { 0, 0, 0, 0 };
static uint32_t rounds = 32;
static unsigned char enc_blob[MAX_LENGTH];
static unsigned char dec_blob[MAX_LENGTH];

void encrypt(uint32_t k[], char *buf, int len, int rounds) __attribute__ ((optimize (0)));
void decrypt(uint32_t k[], char *buf, int len, int rounds) __attribute__ ((optimize (0)));
void scramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d) __attribute__ ((optimize (0)));
void descramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d) __attribute__ ((optimize (0)));

void
scramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d)
{
        uint32_t n, sbox[8];
        uint32_t y = t[0], z = t[1];
        uint32_t delta = DELTA, sum = 0;

        for (n = 0; n < 8; n++) {
                sbox[n] = key[n % 4];
        }
        for (n = 0; n < 32; n++) {
                sum += delta;
                y += ((z << 4) + k[0]) ^ (z+sum) ^ ((z >> 5) + k[1]);
                z += ((y << 4) + k[2]) ^ (y+sum) ^ ((y >> 5) + k[3]);
        }
        t[0] = y; s = d + n; t[1] = z;
}

void
encrypt(uint32_t k[], char *buf, int len, int rounds)
{
        int i;

        if (rounds >= 0) {
                if (rounds % 3 == 0) {
                        k[0] ^= k[1]; k[1] ^= k[0]; k[0] ^= k[1];
                        k[2] ^= k[3]; k[3] ^= k[2]; k[2] ^= k[3];
                } else if (rounds % 3 == 1) {
                        k[0] ^= k[2]; k[2] ^= k[0]; k[0] ^= k[2];
                        k[1] ^= k[3]; k[3] ^= k[1]; k[1] ^= k[3];
                } else if (rounds % 3 == 2) {
                        k[0] ^= k[3]; k[3] ^= k[0]; k[0] ^= k[3];
                        k[2] ^= k[1]; k[1] ^= k[2]; k[2] ^= k[1];
                }
                encrypt(k, buf, len, rounds-1);
        } else {
                for (i = 0; i < len / 8; i++) {
                        scramble(k, (uint32_t *)(buf + (i * 8)), (uint32_t *) 0, (uint32_t *) DELTA);
                }
        }
}

void
descramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d)
{
        uint32_t n, sbox[8];
        uint32_t y = t[0], z = t[1];
        uint32_t delta = DELTA, sum = delta << 5;

        for (n = 0; n < 8; n++) {
                sbox[n] = key[n % 4];
        }
        for (n = 0; n < 32; n++) {
                z -= ((y << 4) + k[2]) ^ (y + sum) ^ ((y >> 5) + k[3]);
                y -= ((z << 4) + k[0]) ^ (z + sum) ^ ((z >> 5) + k[1]);
                sum -= delta;
        }
        t[0] = y; s = d + n; t[1] = z;
}

void
decrypt(uint32_t k[], char *buf, int len, int rounds)
{
        int i;

        if (rounds >= 0) {
                if (rounds % 3 == 0) {
                        k[0] ^= k[1]; k[1] ^= k[0]; k[0] ^= k[1];
                        k[2] ^= k[3]; k[3] ^= k[2]; k[2] ^= k[3];
                } else if (rounds % 3 == 1) {
                        k[0] ^= k[2]; k[2] ^= k[0]; k[0] ^= k[2];
                        k[1] ^= k[3]; k[3] ^= k[1]; k[1] ^= k[3];
                } else if (rounds % 3 == 2) {
                        k[0] ^= k[3]; k[3] ^= k[0]; k[0] ^= k[3];
                        k[2] ^= k[1]; k[1] ^= k[2]; k[2] ^= k[1];
                }
                decrypt(k, buf, len, rounds-1);
        } else {
                for (i = 0; i < len / 8; i++) {
                        descramble(k, (uint32_t *)(buf + (i * 8)), (uint32_t *) 0, (uint32_t *) DELTA);
                }
        }
}

int
key_write(struct file *file, const char __user *ubuf, unsigned long count, void *data)
{
        char buf[MAX_LENGTH];

        printk(KERN_INFO "csaw: called key_write\n");

        memset(buf, 0, sizeof(buf));

        if (count > MAX_LENGTH) {
                count = MAX_LENGTH;
        }
        if (copy_from_user(&buf, ubuf, count)) {
                return -EFAULT;
        }

        sscanf(buf, "%16c\t%d", (char *) key, &rounds);

        return count;
}

int
key_read(char *page, char **start, off_t off, int count, int *eof, void *data)
{
        printk(KERN_INFO "csaw: called key_read\n");

        memcpy(page, key, sizeof(key));

        return sizeof(key);
}

int
encrypt_write(struct file *file, const char __user *ubuf, unsigned long count, void *data)
{
        char buf[MAX_LENGTH];
        uint32_t session[4];

        printk(KERN_INFO "csaw: called encrypt_write\n");

        memset(buf, 0, sizeof(buf));

        if (count > MAX_LENGTH) {
                count = MAX_LENGTH;
        }
        if (copy_from_user(&buf, ubuf, count)) {
                return -EFAULT;
        }

        memcpy(session, key, sizeof(key));

        encrypt(session, buf, sizeof(buf), rounds);

        memcpy(enc_blob, buf, sizeof(buf));

        return count;
}

int
encrypt_read(char *page, char **start, off_t off, int count, int *eof, void *data)
{
        printk(KERN_INFO "csaw: called encrypt_read\n");

        memcpy(page, enc_blob, sizeof(enc_blob));

        return sizeof(enc_blob);
}

int
decrypt_write(struct file *file, const char __user *ubuf, unsigned long count, void *data)
{
        char buf[MAX_LENGTH];
        uint32_t session[4];

        printk(KERN_INFO "csaw: called decrypt_write\n");

        memset(buf, 0, sizeof(buf));

        if (count > MAX_LENGTH) {
                count = MAX_LENGTH;
        }
        if (copy_from_user(&buf, ubuf, count)) {
                return -EFAULT;
        }

        memcpy(session, key, sizeof(key));

        decrypt(session, buf, sizeof(buf), rounds);

        memcpy(dec_blob, buf, sizeof(buf));

        return count;
}

int
decrypt_read(char *page, char **start, off_t off, int count, int *eof, void *data)
{
        printk(KERN_INFO "csaw: called decrypt_read\n");

        memcpy(page, dec_blob, sizeof(dec_blob));

        return sizeof(dec_blob);
}

int __init
csaw_init(void)
{
        printk(KERN_INFO "csaw: loading module\n");

        memset(enc_blob, 0, sizeof(enc_blob));
        memset(dec_blob, 0, sizeof(dec_blob));

        proc_csaw = proc_mkdir(PROC_CSAW, NULL);

        proc_key = create_proc_entry(PROC_KEY, 0666, proc_csaw);
        proc_key->read_proc = key_read;
        proc_key->write_proc = key_write;

        proc_encrypt = create_proc_entry(PROC_ENCRYPT, 0666, proc_csaw);
        proc_encrypt->read_proc = encrypt_read;
        proc_encrypt->write_proc = encrypt_write;

        proc_decrypt = create_proc_entry(PROC_DECRYPT, 0666, proc_csaw);
        proc_decrypt->read_proc = decrypt_read;
        proc_decrypt->write_proc = decrypt_write;

        printk(KERN_INFO "csaw: created /proc/csaw2011 entries\n");

        return 0;
}

void __exit
csaw_exit(void)
{
        printk(KERN_INFO "csaw: unloading module\n");

        remove_proc_entry(PROC_KEY, proc_csaw);
        remove_proc_entry(PROC_ENCRYPT, proc_csaw);
        remove_proc_entry(PROC_DECRYPT, proc_csaw);
        remove_proc_entry(PROC_CSAW, NULL);

        printk(KERN_INFO "csaw: removed /proc/csaw2011 entries\n");
}

module_init(csaw_init);
module_exit(csaw_exit);
